import networkx as nx
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np

mpl.use("TkAgg")

pos_default = {1: np.array([-0.43761481, 0.77035406]), 2: np.array([0.54153089, 0.53466979]),
               3: np.array([0.12276006, 0.37285459]), 4: np.array([-0.05237984, -0.0374406]),
               5: np.array([-0.05544236, 0.73358622]), 6: np.array([0.2440427, -0.02875812]),
               7: np.array([0.68812098, 0.03353095]), 8: np.array([-0.48913981, -0.29508746]),
               9: np.array([-0.37000831, 0.08967374]), 10: np.array([0.16850579, -0.969913]),
               11: np.array([-0.62162977, 0.23964767]), 12: np.array([0.33987234, -0.47598375]),
               13: np.array([-0.06168556, -0.5953669]), 14: np.array([-0.45019673, -0.83994827]),
               15: np.array([0.76653106, -0.247743]), 16: np.array([0.48167638, 1.]),
               17: np.array([-0.81494301, -0.28407591])}


def draw_graph(edge_values: np.ndarray, node_values: np.ndarray, filepath: str, node_labels=None, show=False,
               random_pos=False, use_node_values=True):
    if not use_node_values:
        plt.figure(figsize=(10, 10))
        node_values = np.ones(edge_values.shape[0])
        font_color = "khaki"
        edge_color = "dimgrey"
    else:
        plt.figure(figsize=(13, 10))
        font_color = "maroon"
        edge_color = None
    node_i, node_j = np.meshgrid(np.arange(1, node_values.shape[0] + 1), np.arange(1, node_values.shape[0] + 1),
                                 indexing="ij")
    graph = nx.Graph()


    cmap = plt.cm.autumn_r

    def layout(g):
        return nx.spring_layout(g, seed=13648)

    indexes = np.stack((node_i, node_j, edge_values), axis=-1).reshape(-1, 3)
    indexes = [(int(u), int(v), w)
               for [u, v, w] in indexes.tolist() if int(u) != int(v)]
    graph.add_weighted_edges_from(indexes)
    if random_pos:
        pos = layout(graph)
    else:
        pos = pos_default
    print(pos)
    edges = np.array(graph.edges())
    edge_width = edge_values[edges[:, 0] - 1, edges[:, 1] - 1]

    color_normal = mpl.colors.Normalize(
        vmin=min(node_values), vmax=max(node_values))
    if not use_node_values:
        node_colors = "cornflowerblue"
    else:
        node_colors = [cmap(color_normal(val)) for val in node_values]
    if np.max(node_values) == np.min(node_values):
        node_values = np.ones_like(node_values, dtype=np.float32) * 0.5
    else:
        node_values = (node_values - np.min(node_values)) / (np.max(node_values) - np.min(node_values))
    nx.draw_networkx_nodes(graph,
                           pos,
                           node_size=node_values * 1900 + 600,
                           node_color=node_colors,
                           linewidths=1.0,
                           edgecolors=edge_color)
    edges = nx.draw_networkx_edges(graph,
                                   pos,
                                   width=edge_width * 2.5,
                                   node_size=node_values,
                                   edge_color="dimgrey",
                                   # edge_cmap=cmap,
                                   # edge_vmin=min(edge_width),
                                   # edge_vmax=max(edge_width),
                                   # style='dashed',
                                   alpha=1)
    if node_labels is not None:
        nx.draw_networkx_labels(
            graph, pos, node_labels, font_size=15, font_color=font_color, font_weight='bold')
    else:
        nx.draw_networkx_labels(graph, pos, dict([(i, f"{i}") for i in range(1, len(node_values) + 1)]), font_size=15,
                                font_color=font_color, font_weight='bold')
    edges.set_array(edge_width)
    ax = plt.gca()
    ax.set_axis_off()
    if use_node_values:
        plt.colorbar(edges, ax=ax)

    plt.savefig(filepath)
    if show:
        plt.show()


def example():
    node_values = np.random.rand(17)
    edge_values = np.random.rand(17, 17)
    draw_graph(edge_values, node_values, "../pictures/graph1.png", dict([(i, f"{i}") for i in range(1, 18)]),
               use_node_values=False)


if __name__ == '__main__':
    example()
